import requests
import json

# OpenAI 兼容接口
response = requests.post(
    "http://localhost:11434/v1/chat/completions",
    json={
        "model": "deepseek-r1:1.5b",
        "messages": [
            {
                "role": "user",
                "content": "如何在月球建立一个基地，拆分成为100个环节，环节比如：从构建地球建设基地，到交通运输，卫星建设，发送卫星，月球落地等，一共细分为100个。并对每个环节进行分析，论证，规划，设计，并考虑到最基础的数学和物理学公式。需要特别特别详细的设计和分析。最后，在结尾添加设计5个小球在月亮上飞的游戏，并显示所有代码.这样的游戏设计10个，都需要代码"
            }
        ],
        "stream": True,
        "temperature": 0.1,
        "max_tokens": 8192,
    }
)

total_tokens = 0
generated_text = ""

# 逐块处理流式响应
for line in response.iter_lines():
    if line:
        chunk = line.decode("utf-8")
        if chunk.startswith("data: "):
            try:
                # 解析 JSON 数据
                data = json.loads(chunk[5:])  # 去除 "data: " 前缀
                if "choices" in data:
                    choice = data["choices"][0]["delta"]
                    # 累积生成的文本
                    generated_text += choice["content"]
                    # 计算当前块的 Token 数（假设文本的长度近似 Token 数）
                    current_tokens = len(choice["content"])
                    total_tokens += current_tokens
                    print(f"新生成内容: {choice['content'].strip()}")
                    print(f"当前累计 Token: {total_tokens}")
                elif "error" in data:
                    print(f"错误: {data['error']}")
            except json.JSONDecodeError as e:
                print(f"解析错误: {e}")

print("\n最终生成文本:")
print(generated_text)
print(f"总 Token 数: {total_tokens}")